/*
 * Javassist, a Java-bytecode translator toolkit.
 * Copyright (C) 1999- Shigeru Chiba. All Rights Reserved.
 *
 * The contents of this file are subject to the Mozilla Public License Version
 * 1.1 (the "License"); you may not use this file except in compliance with
 * the License.  Alternatively, the contents of this file may be used under
 * the terms of the GNU Lesser General Public License Version 2.1 or later,
 * or the Apache License Version 2.0.
 *
 * Software distributed under the License is distributed on an "AS IS" basis,
 * WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
 * for the specific language governing rights and limitations under the
 * License.
 */

package com.feilong.lib.javassist.util.proxy;

import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.reflect.Method;
import java.security.ProtectionDomain;
import java.util.List;

import com.feilong.lib.javassist.CannotCompileException;
import com.feilong.lib.javassist.bytecode.ClassFile;

/**
 * Helper class for invoking {@link ClassLoader#defineClass(String,byte[],int,int)}.
 *
 * @since 3.22
 */
public class DefineClassHelper{

    private static abstract class Helper{

        abstract Class<?> defineClass(
                        String name,
                        byte[] b,
                        int off,
                        int len,
                        Class<?> neighbor,
                        ClassLoader loader,
                        ProtectionDomain protectionDomain) throws ClassFormatError,CannotCompileException;
    }

    private static class Java9 extends Helper{

        final class ReferencedUnsafe{

            private final SecurityActions.TheUnsafe sunMiscUnsafeTheUnsafe;

            private final MethodHandle              defineClass;

            ReferencedUnsafe(SecurityActions.TheUnsafe usf, MethodHandle meth){
                this.sunMiscUnsafeTheUnsafe = usf;
                this.defineClass = meth;
            }

            Class<?> defineClass(String name,byte[] b,int off,int len,ClassLoader loader,ProtectionDomain protectionDomain)
                            throws ClassFormatError{
                try{
                    if (getCallerClass.invoke(stack) != Java9.class){
                        throw new IllegalAccessError("Access denied for caller.");
                    }
                }catch (Exception e){
                    throw new RuntimeException("cannot initialize", e);
                }
                try{
                    return (Class<?>) defineClass
                                    .invokeWithArguments(sunMiscUnsafeTheUnsafe.theUnsafe, name, b, off, len, loader, protectionDomain);
                }catch (Throwable e){
                    if (e instanceof RuntimeException){
                        throw (RuntimeException) e;
                    }
                    if (e instanceof ClassFormatError){
                        throw (ClassFormatError) e;
                    }
                    throw new ClassFormatError(e.getMessage());
                }
            }
        }

        private final Object           stack;

        private final Method           getCallerClass;

        private final ReferencedUnsafe sunMiscUnsafe = getReferencedUnsafe();

        Java9(){
            Class<?> stackWalkerClass = null;
            try{
                stackWalkerClass = Class.forName("java.lang.StackWalker");
            }catch (ClassNotFoundException e){
                // Skip initialization when the class doesn't exist i.e. we are on JDK < 9
            }
            if (stackWalkerClass != null){
                try{
                    Class<?> optionClass = Class.forName("java.lang.StackWalker$Option");
                    stack = stackWalkerClass.getMethod("getInstance", optionClass)
                                    // The first one is RETAIN_CLASS_REFERENCE
                                    .invoke(null, optionClass.getEnumConstants()[0]);
                    getCallerClass = stackWalkerClass.getMethod("getCallerClass");
                }catch (Throwable e){
                    throw new RuntimeException("cannot initialize", e);
                }
            }else{
                stack = null;
                getCallerClass = null;
            }
        }

        private final ReferencedUnsafe getReferencedUnsafe(){
            try{
                if (privileged != null && getCallerClass.invoke(stack) != this.getClass()){
                    throw new IllegalAccessError("Access denied for caller.");
                }
            }catch (Exception e){
                throw new RuntimeException("cannot initialize", e);
            }
            try{
                SecurityActions.TheUnsafe usf = SecurityActions.getSunMiscUnsafeAnonymously();
                List<Method> defineClassMethod = usf.methods.get("defineClass");
                // On Java 11+ the defineClass method does not exist anymore
                if (null == defineClassMethod){
                    return null;
                }
                MethodHandle meth = MethodHandles.lookup().unreflect(defineClassMethod.get(0));
                return new ReferencedUnsafe(usf, meth);
            }catch (Throwable e){
                throw new RuntimeException("cannot initialize", e);
            }
        }

        @Override
        Class<?> defineClass(String name,byte[] b,int off,int len,Class<?> neighbor,ClassLoader loader,ProtectionDomain protectionDomain)
                        throws ClassFormatError{
            try{
                if (getCallerClass.invoke(stack) != DefineClassHelper.class){
                    throw new IllegalAccessError("Access denied for caller.");
                }
            }catch (Exception e){
                throw new RuntimeException("cannot initialize", e);
            }
            return sunMiscUnsafe.defineClass(name, b, off, len, loader, protectionDomain);
        }
    }

    private static class Java7 extends Helper{

        private final SecurityActions stack       = SecurityActions.stack;

        private final MethodHandle    defineClass = getDefineClassMethodHandle();

        private final MethodHandle getDefineClassMethodHandle(){
            if (privileged != null && stack.getCallerClass() != this.getClass()){
                throw new IllegalAccessError("Access denied for caller.");
            }
            try{
                return SecurityActions.getMethodHandle(
                                ClassLoader.class,
                                "defineClass",
                                new Class[] { String.class, byte[].class, int.class, int.class, ProtectionDomain.class });
            }catch (NoSuchMethodException e){
                throw new RuntimeException("cannot initialize", e);
            }
        }

        @Override
        Class<?> defineClass(String name,byte[] b,int off,int len,Class<?> neighbor,ClassLoader loader,ProtectionDomain protectionDomain)
                        throws ClassFormatError{
            if (stack.getCallerClass() != DefineClassHelper.class){
                throw new IllegalAccessError("Access denied for caller.");
            }
            try{
                return (Class<?>) defineClass.invokeWithArguments(loader, name, b, off, len, protectionDomain);
            }catch (Throwable e){
                if (e instanceof RuntimeException){
                    throw (RuntimeException) e;
                }
                if (e instanceof ClassFormatError){
                    throw (ClassFormatError) e;
                }
                throw new ClassFormatError(e.getMessage());
            }
        }
    }

    private static class JavaOther extends Helper{

        private final Method          defineClass = getDefineClassMethod();

        private final SecurityActions stack       = SecurityActions.stack;

        private final Method getDefineClassMethod(){
            if (privileged != null && stack.getCallerClass() != this.getClass()){
                throw new IllegalAccessError("Access denied for caller.");
            }
            try{
                return SecurityActions.getDeclaredMethod(
                                ClassLoader.class,
                                "defineClass",
                                new Class[] { String.class, byte[].class, int.class, int.class, ProtectionDomain.class });
            }catch (NoSuchMethodException e){
                throw new RuntimeException("cannot initialize", e);
            }
        }

        @Override
        Class<?> defineClass(String name,byte[] b,int off,int len,Class<?> neighbor,ClassLoader loader,ProtectionDomain protectionDomain)
                        throws ClassFormatError,CannotCompileException{
            Class<?> klass = stack.getCallerClass();
            if (klass != DefineClassHelper.class && klass != this.getClass()){
                throw new IllegalAccessError("Access denied for caller.");
            }
            try{
                SecurityActions.setAccessible(defineClass, true);
                return (Class<?>) defineClass.invoke(loader, new Object[] { name, b, off, len, protectionDomain });
            }catch (Throwable e){
                if (e instanceof ClassFormatError){
                    throw (ClassFormatError) e;
                }
                if (e instanceof RuntimeException){
                    throw (RuntimeException) e;
                }
                throw new CannotCompileException(e);
            }finally{
                SecurityActions.setAccessible(defineClass, false);
            }
        }
    }

    // Java 11+ removed sun.misc.Unsafe.defineClass, so we fallback to invoking defineClass on
    // ClassLoader until we have an implementation that uses MethodHandles.Lookup.defineClass
    private static final Helper privileged = ClassFile.MAJOR_VERSION >= ClassFile.JAVA_9 ? new Java9()
                    : ClassFile.MAJOR_VERSION >= ClassFile.JAVA_7 ? new Java7() : new JavaOther();

    /**
     * Loads a class file by a given class loader.
     *
     * <p>
     * This first tries to use {@code java.lang.invoke.MethodHandle} to load a class.
     * Otherwise, or if {@code neighbor} is null,
     * this tries to use {@code sun.misc.Unsafe} to load a class.
     * Then it tries to use a {@code protected} method in {@code java.lang.ClassLoader}
     * via {@code PrivilegedAction}. Since the latter approach is not available
     * any longer by default in Java 9 or later, the JVM argument
     * {@code --add-opens java.base/java.lang=ALL-UNNAMED} must be given to the JVM.
     * If this JVM argument cannot be given, {@link #toPublicClass(String,byte[])}
     * should be used instead.
     * </p>
     *
     * @param className
     *            the name of the loaded class.
     * @param neighbor
     *            the class contained in the same package as the loaded class.
     * @param loader
     *            the class loader. It can be null if {@code neighbor} is not null
     *            and the JVM is Java 11 or later.
     * @param domain
     *            if it is null, a default domain is used.
     * @param bcode
     *            the bytecode for the loaded class.
     * @since 3.22
     */
    public static Class<?> toClass(String className,Class<?> neighbor,ClassLoader loader,ProtectionDomain domain,byte[] bcode)
                    throws CannotCompileException{
        try{
            return privileged.defineClass(className, bcode, 0, bcode.length, neighbor, loader, domain);
        }catch (RuntimeException e){
            throw e;
        }catch (CannotCompileException e){
            throw e;
        }catch (ClassFormatError e){
            Throwable t = e.getCause();
            throw new CannotCompileException(t == null ? e : t);
        }catch (Exception e){
            throw new CannotCompileException(e);
        }
    }

    private DefineClassHelper(){
    }
}
